import os
import numpy as np
import mediapipe as mp
from mediapipe.framework.formats import landmark_pb2

from pose_labeling_utils import PoseLabelingUtils


class MediaPipeUtils:
    def __init__(self, model_path, running_mode, num_poses):
        self.mp_drawing = mp.solutions.drawing_utils
        self.mp_pose = mp.solutions.pose
        self.mp_drawing_styles = mp.solutions.drawing_styles
        self.pl_utils = PoseLabelingUtils()

        self.detector = self.initialize_mediapipe(
            model_path=model_path, running_mode=running_mode,
            num_poses=num_poses, min_pose_detection_confidence=0.7,
            min_pose_presence_confidence=0.7, min_tracking_confidence=0.7)

    def initialize_mediapipe(self, model_path, running_mode, num_poses,
                             min_pose_detection_confidence,
                             min_pose_presence_confidence,
                             min_tracking_confidence):
        BaseOptions = mp.tasks.BaseOptions
        PoseLandmarker = mp.tasks.vision.PoseLandmarker
        PoseLandmarkerOptions = mp.tasks.vision.PoseLandmarkerOptions
        VisionRunningMode = mp.tasks.vision.RunningMode

        options = PoseLandmarkerOptions(
            base_options=BaseOptions(model_asset_path=model_path),
            running_mode=getattr(VisionRunningMode, running_mode),
            num_poses=num_poses,
            min_pose_detection_confidence=min_pose_detection_confidence,
            min_pose_presence_confidence=min_pose_presence_confidence,
            min_tracking_confidence=min_tracking_confidence,
            output_segmentation_masks=True
        )
        return PoseLandmarker.create_from_options(options)

    def detect(self, frame):
        # 将图像转换为 MediaPipe 图像对象
        img = mp.Image(image_format=mp.ImageFormat.SRGB, data=frame)
        return self.detector.detect(frame)

    def detect_for_video(self, frame, timestamp_ms):
        # 将视频转换为 MediaPipe 图像对象
        img = mp.Image(image_format=mp.ImageFormat.SRGB, data=frame)
        return self.detector.detect_for_video(img, timestamp_ms)

    def draw_landmarks_on_image(self, rgb_image, detection_result):
        pose_landmarks_list = detection_result.pose_landmarks
        annotated_image = np.copy(rgb_image)

        for idx in range(len(pose_landmarks_list)):
            pose_landmarks = pose_landmarks_list[idx]

            # Draw the pose landmarks.
            pose_landmarks_proto = landmark_pb2.NormalizedLandmarkList()
            pose_landmarks_proto.landmark.extend([
                landmark_pb2.NormalizedLandmark(x=landmark.x, y=landmark.y, z=landmark.z) for landmark in pose_landmarks
            ])

            self.mp_drawing.draw_landmarks(
                annotated_image,
                pose_landmarks_proto,
                self.mp_pose.POSE_CONNECTIONS,
                self.mp_drawing_styles.get_default_pose_landmarks_style()
            )

            # 提取关键点
            file_path = '../data/keypoints_and_thresholds.json'
            keypoints = self.get_existing_keypoints(file_path, pose_landmarks)

            # 计算头部姿势
            head_labels = self.pl_utils.label_head_pose(keypoints)
            # 计算臂部姿势
            arm_labels = self.pl_utils.label_arm_pose(keypoints)
            # 计算腿部姿势
            leg_labels = self.pl_utils.label_leg_pose(keypoints)

            # 在图像上绘制动作标签
            action_label = f"{head_labels}, {arm_labels}, {leg_labels}"
            nose = pose_landmarks[self.mp_pose.PoseLandmark.NOSE.value]
            label_position = (int(nose.x * rgb_image.shape[1]), int(nose.y * rgb_image.shape[0]))
            annotated_image = self.pl_utils.draw_action_label(
                annotated_image, action_label, label_position)

        return annotated_image

    def lv_to_np(self, pose_landmarks, key, need_v=False):
        value = pose_landmarks[self.mp_pose.PoseLandmark[key].value]
        x = value.x
        y = value.y
        visibility = value.visibility
        keypoint = np.array([x, y, visibility])
        if not need_v:
            keypoint = np.array([x, y])
            return keypoint
        return keypoint

    def get_existing_keypoints(self, file_path, pose_landmarks):
        pose_landmark_index = self.pl_utils.load_keypoints_and_thresholds(file_path)[1]
        keypoints = {}
        for k, v in pose_landmark_index.items():
            kp = self.lv_to_np(pose_landmarks, k)
            keypoints[k.lower()] = kp
        return keypoints
